from __future__ import annotations
from msgspec import Struct, field, json
from typing import Literal, Union, TypeAlias
from src.dto import MessageStruct
import uuid
import secrets
import struct

class TypedStruct(Struct):
    type: str

OutType: TypeAlias = Union["HelloMessage", "ListenMessage", "DetectMessage", dict]

class MessageDecoder:

    @classmethod
    def from_bytes(cls, data: bytes) -> OutType:
        try:
            typed_struct = json.Decoder(TypedStruct).decode(data)
        except Exception as e:
            raise ValueError(f"Failed to decode message {data}: {e}")
        match typed_struct.type:
            case "hello":
                return HelloMessage.from_bytes(data)
            case "listen":
                pass
            case "detect":
                pass
            case _:
                # 自适应解析
                return json.decode(data)

class AudioParams(MessageStruct):
    """音频参数配置"""
    format: str = field(default_factory=lambda:"opus")
    sample_rate: int = field(default_factory=lambda:16000)
    channels: int = field(default_factory=lambda:1)
    frame_duration: int = field(default_factory=lambda:20)


class Features(MessageStruct):
    """功能特性配置"""
    aec: bool = field(default_factory=lambda:False)  # 回声消除
    mcp: bool = field(default_factory=lambda:False)  # MCP 协议支持


class HelloMessage(MessageStruct):
    """Hello 消息结构"""
    type: str = field(default_factory=lambda:"hello")
    version: int = field(default_factory=lambda:3)  # 版本号(WS为1, MQTT为3)
    features: Features = field(default_factory=Features)
    transport: Literal["websocket", "udp"] = field(default_factory=lambda:"websocket")
    audio_params: AudioParams = field(default_factory=AudioParams)


class UdpInfo(MessageStruct):
    """UDP连接信息"""
    @staticmethod
    def _generate_udp_encryption():
        """生成UDP加密算法，固定为aes-128-ctr"""
        return "aes-128-ctr"

    @staticmethod
    def _generate_udp_key():
        """生成16字节随机密钥，转换为十六进制字符串"""
        return secrets.token_bytes(16).hex()

    @staticmethod
    def _generate_udp_nonce():
        """生成UDP nonce，模拟app.js中的generateUdpHeader(0, 0, 0)"""
        # 根据app.js的generateUdpHeader实现:
        # buffer.writeUInt8(1, 0);        // type: 1
        # buffer.writeUInt16BE(0, 2);     // length: 0  
        # buffer.writeUInt32BE(0, 4);     // connectionId: 0
        # buffer.writeUInt32BE(0, 8);     // timestamp: 0
        # buffer.writeUInt32BE(0, 12);    // sequence: 0
        header = struct.pack('>BxHLLL', 1, 0, 0, 0, 0)  # 16字节头部
        return header.hex()
    
    server: str = field(default_factory=lambda:"192.168.2.1")  # UDP服务器地址
    port: int = field(default_factory=lambda:8888) # UDP端口
    encryption: str = field(default_factory=_generate_udp_encryption)  # UDP加密算法
    key: str = field(default_factory=_generate_udp_key)  # UDP密钥(十六进制字符串)
    nonce: str = field(default_factory=_generate_udp_nonce)  # UDP随机数(十六进制字符串)



class HelloResponse(MessageStruct):
    """Hello 响应消息结构"""
    type: str = field(default_factory=lambda: "hello")
    session_id: str = field(default_factory=lambda:str(uuid.uuid4())[:8])  # 会话ID
    udp: UdpInfo = field(default_factory=UdpInfo) # UDP连接信息
    transport: Literal["websocket", "udp"] = field(default_factory=lambda: "udp")  # 传输类型
    version: int = field(default_factory=lambda:3)  # 版本号(MQTT为3)

class CustomMessage(MessageStruct):
    """自定义消息结构"""
    type: str = field(default_factory=lambda: "custom")
    payload: dict = field(default_factory=dict)  # 自定义载荷数据，对应C++中的payload对象

class ImageMessage(MessageStruct):
    """图片消息结构"""
    url: str
    width: int = field(default_factory=lambda: 128)
    height: int = field(default_factory=lambda: 64)
    format: Literal["jpeg", "png", "bmp", "rgb565"] = field(default_factory=lambda: "rgb565")
    type: str = field(default_factory=lambda: "image")

class ListenMessage(MessageStruct):
    """Listen 消息结构"""
    type: str = field(default_factory=lambda: "listen")

class DetectMessage(MessageStruct):
    """Detect 消息结构"""
    type: str = field(default_factory=lambda: "detect")